#!/usr/bin/env python
# -*- coding:utf-8 -*-
import socket,struct
from libs import printer
protocols = {
    "file": "file://{}",
    "http": "http://{}:{}/{}",
    "httpauth": "http://{}:{}@{}:{}/{}",
    "https": "https://{}:{}/{}",
    "dict": "dict://{}:{}/{}",
    "sftp": "sftp://{}:{}/{}",
    "tftp": "tftp://{}:{}/{}",
    "ldap": "ldap://{}:{}/{}",
    "gopher": "gopher://{}:{}/_{}"
}

def wrapper_protocol(proto:str):
    result = ""
    if proto == "file":
        pdata = input("请输入file协议数据: ")
        result = protocols[proto].format(pdata.strip())
    elif proto == "http":
        pdata = input("请输入http协议数据[ip@port@data]: ")
        ip,port,data = pdata.strip().split("@")
        result = protocols[proto].format(ip,port,data)
    elif proto == "httpauth":
        pdata = input("请输入httpauth协议数据[username@password@ip@port@data]: ")
        u,p,ip,port,data = pdata.strip().split("@")
        result = protocols[proto].format(u,p,ip,port,data)
    elif proto == "https":
        pdata = input("请输入https协议数据[ip@port@data]: ")
        ip,port,data = pdata.strip().split("@")
        result = protocols[proto].format(ip,port,data)
    elif proto == "dict":
        pdata = input("请输入dict协议数据[ip@port@data]: ")
        ip,port,data = pdata.strip().split("@")
        result = protocols[proto].format(ip,port,data)
    elif proto == "sftp":
        pdata = input("请输入sftp协议数据[ip@port@data]: ")
        ip,port,data = pdata.strip().split("@")
        result = protocols[proto].format(ip,port,data)
    elif proto == "tftp":
        pdata = input("请输入tftp协议数据[ip@port@data]: ")
        ip,port,data = pdata.strip().split("@")
        result = protocols[proto].format(ip,port,data)
    elif proto == "ldap":
        pdata = input("请输入ldap协议数据[ip@port@data]: ")
        ip,port,data = pdata.strip().split("@")
        result = protocols[proto].format(ip,port,data)
    elif proto == "gopher":
        pdata = input("请输入gopher协议数据[ip@port@data]: ")
        ip,port,data = pdata.strip().split("@")
        result = protocols[proto].format(ip,port,data)
    else:
        pass
    return result
def ip_to_decimal(ip:str):
    result = ""
    wraip = socket.inet_aton(ip)
    result = struct.unpack("!l", wraip)[0]
    return result

def ip_to_ota(ip:str):
    result = ""
    result = ".".join([str(oct(int(p))).replace("o","") for p in ip.split(".")])
    return result
def show_menu():
    printer.info("SSRF选项: ")
    printer.plus("1. 协议包装")
    printer.plus("2. IP进制转换")
def ip_menu():
    printer.info("IP转换: ")
    printer.plus("1. 十进制")
    printer.plus("2. 八进制")    
def execute(args):
    try:
        show_menu()
        choice = input("请输入选项: ")
        if choice == "1":
            proto = input("请输入协议名称: ")
            proto = proto.strip()
            if proto in protocols.keys():
                result = wrapper_protocol(proto)
                printer.info("包装结果: " + result)
            else:
                printer.warn("协议不存在,可自行增加!")
        elif choice == "2":
            ip_menu()
            c = input("[IP]请输入选项: ")
            if c == "1":
                ip = input("请输入IP: ")
                ip = ip.strip()
                result = ip_to_decimal(ip)
                printer.info("转换结果: " + str(result))
            elif c == "2":
                ip = input("请输入IP: ")
                ip = ip.strip()
                result = ip_to_ota(ip)
                printer.info("转换结果: " + result)
            else:
                printer.warn("选项不存在.")
        else:
            printer.warn("选项不存在.")             
    except Exception as e:
        printer.warn("ssrf程序错误!")
        print(e)
